import pandas as pd
from sklearn.model_selection import train_test_split

# 读取CSV文件
df = pd.read_csv('cs_data.csv')

# 将数据分为训练集和临时集（测试集+验证集）
# test_size=0.3 意味着30%的数据将进入临时集，70%进入训练集
# random_state=42 确保结果可重现
train_df, temp_df = train_test_split(df, test_size=0.2, random_state=42)

# 将临时集再分为测试集和验证集
# test_size=0.5 意味着临时集的一半将成为测试集，另一半成为验证集
test_df, val_df = train_test_split(temp_df, test_size=0.5, random_state=42)

# 打印每个集合的大小
print(f"训练集大小: {len(train_df)}")
print(f"测试集大小: {len(test_df)}")
print(f"验证集大小: {len(val_df)}")

# 保存分割后的数据集
# index=False 确保不将索引列写入CSV文件
train_df.to_csv('train_data.csv', index=False)
test_df.to_csv('test_data.csv', index=False)
val_df.to_csv('val_data.csv', index=False)

print("数据集已成功分割并保存。")